{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MLRun secret handling using HashiCorp Vault and Azure key-store\n",
    "\n",
    "This notebook demonstrates secrets creation and handling in MLRun, both when using Vault and using Azure key-store.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Create and deploy a function displaying secret values\n",
    "\n",
    "For this demo we'll use a simple function which gets a list of secret names to show, attempts to get their value from the context, and prints their contents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlrun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mlrun: start-code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vault_func(context, secrets: list):\n",
    "    \"\"\"Validate that given secrets exists\n",
    "\n",
    "    :param context: the MLRun context\n",
    "    :param secrets: name of the secrets that we want to look at\n",
    "    \"\"\"\n",
    "    context.logger.info(\"running function\")\n",
    "    for sec_name in secrets:\n",
    "        sec_value = context.get_secret(sec_name)\n",
    "        context.logger.info(\"Secret name: {}, value: {}\".format(sec_name, sec_value))\n",
    "\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mlrun: end-code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "func = mlrun.code_to_function(name=\"vault-func\", kind=\"job\", image=\"mlrun/mlrun\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with secrets kept in HashiCorp Vault"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a project & initialize Vault support\n",
    "When a project is created, the `create_vault_secrets` command has to be used to request that the underlying framework is created that will enable Vault secrets to be used with this project. Calling this method on the project will create the following constructs for the project (if not already existing):\n",
    "\n",
    "1. A k8s serviceaccount (`sa-vault-{project name}`)\n",
    "2. A Vault policy (`mlrun-project-{project name}`) that enables access to secrets in the project path (`/secrets/secret/mlrun/projects/{proj name}`)\n",
    "3. A Vault k8s role (`mlrun-role-project-{project name}`) that associates the SA's token with the policy\n",
    "\n",
    "These configurations are performed on the MLRun API server side, not from the client.\n",
    "In addition, the command passes a list of secret values to be associated with the project in Vault.\n",
    "\n",
    "When a project was already initialized in Vault, those secrets can be examined using the `get_vault_secrets` method.\n",
    "This method runs on the API server but used the client's Vault token, so client needs to be authorized to access Vault to execute it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_name = \"vault-mlrun\"\n",
    "\n",
    "proj = mlrun.new_project(proj_name)\n",
    "\n",
    "project_secrets = {\"aws_key\": \"1234567890\", \"github_key\": \"proj1Key!!!\"}\n",
    "proj.create_vault_secrets(project_secrets)\n",
    "\n",
    "proj.get_vault_secrets()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Initialize function runtime\n",
    "The `.with_secrets` function has a '`vault`' secret kind that will pass the specified Vault project secrets to the function context. The function spec\n",
    "only contains the keys of the secrets ('aws_key' etc.) - the actual secret value is retrieved from Vault and planted in the function \n",
    "context in runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = mlrun.new_task(\n",
    "    project=proj_name,\n",
    "    name=\"vault_test_run\",\n",
    "    handler=\"vault_func\",\n",
    "    params={\"secrets\": [\"github_key\", \"aws_key\"]},\n",
    ")\n",
    "\n",
    "# Add access to project-level secrets\n",
    "task.with_secrets(\"vault\", [\"aws_key\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute the function\n",
    "Running the function using the task we've created will show only the value of the `aws_key` secret, since it's the only secret passed to the `task` object. Modifying the `with_secrets` command will result in different secrets being available to the function runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = func.run(task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By modifying the call to `with_secrets` so that all project secrets are passed to the runtime, the function will be able to present both the `github_key` and the `aws_key` secret values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Access to all project-level secrets can be obtained by passing an empty list of secret names\n",
    "task.with_secrets(\"vault\", [])\n",
    "\n",
    "result = func.run(task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Run the same function in another project's context\n",
    "We will create a 2nd project, and assign different secret values to it. When the same function is executed in the new project's runtime context, it will get \n",
    "the new project's secrets. When running in this context, the function has no access to other projects' secrets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_name_2 = \"vault-mlrun-2\"\n",
    "proj2 = mlrun.new_project(proj_name_2)\n",
    "proj2.create_vault_secrets(\n",
    "    {\"aws_key\": \"0987654321\", \"github_key\": \"proj2Key???\", \"password\": \"myPassword\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task2 = mlrun.new_task(\n",
    "    project=proj_name_2,\n",
    "    name=\"vault_test_run_2\",\n",
    "    handler=\"vault_func\",\n",
    "    params={\"secrets\": [\"password\", \"github_key\", \"aws_key\"]},\n",
    ")\n",
    "task2.with_secrets(\"vault\", [\"aws_key\", \"github_key\", \"password\"])\n",
    "\n",
    "result = func.run(task2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Project-level client side Vault operations\n",
    "Client-side Vault integration is possible, assuming that the pod hosting this notebook has Vault connectivity and is properly configured. \n",
    "Vault integration can be done using the `.with_secrets()` function, and then by using the `.get_secret()` utility function to extract secret value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "proj.with_secrets(\"vault\", [\"github_key\"])\n",
    "proj.get_secret(\"github_key\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with secrets kept in Azure key vault"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration\n",
    "To work with Azure key vault, you need to first setup the following:\n",
    "1. Setup a key vault in your Azure subscription. For the purposes of this demo we'll assume the key vault is called \"azure-demo-key-vault\"\n",
    "2. Create a service principal in Azure that will be granted access to the key vault. For creating a service principal through the Azure portal follow the steps in this page: https://docs.microsoft.com/en-us/azure/active-directory/develop/howto-create-service-principal-portal\n",
    "3. Assign a key vault access policy to the service principal, as described in https://docs.microsoft.com/en-us/azure/key-vault/general/assign-access-policy-portal\n",
    "4. Create a secret access key for the service principal, following the steps in: https://docs.microsoft.com/en-us/azure/active-directory/develop/howto-create-service-principal-portal#get-tenant-and-app-id-values-for-signing-in. Make sure you have access to the following 3 identifiers:\n",
    "    1. Directory (tenant) id\n",
    "    2. Application (client) id\n",
    "    3. Secret key\n",
    "5. Generate a new k8s secret containing the above mentioned 3 keys. This can be done using `kubectl` with a command such as the following:\n",
    "\n",
    "    >``` bash\n",
    "    > kubectl -n <namespace> create secret generic azure-key-vault-secret \\\n",
    "    >    --from-literal=secret=<secret key> \\\n",
    "    >    --from-literal=tenant_id=<tenant id> \\\n",
    "    >    --from-literal=client_id=<client id>\n",
    "    >```\n",
    "    \n",
    "    The names of the values within the secret are important, make sure you don't change them, as the MLRun code expects them to have these specific names.\n",
    "6. Populate the key vault with secrets you wish your code to access. In this demo we assume two secrets are created - `demo-secret-1` and `demo-secret-2`.\n",
    "    \n",
    "    > *Note:*\n",
    "    > Azure key vault supports 3 types of entities - keys, secrets and certificates. The MLRun code supports accessing only secrets.\n",
    "\n",
    "Every pod that needs access to Azure key vault secrets must have this secret mounted. MLRun will automatically mount this secret on execution pods that are generated through its `run` command. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Executing code that accesses Azure key vault secrets\n",
    "To access secrets, the `.with_secrets()` function should be called, specifying a secret source of type `azure_vault`. The name of the key vault and the name of the k8s secret (created in previous section) needs to be specified as part of the secret source definition. The following code generates a new execution task and configures it to access the Azure key vault called `azure-demo-key-vault`, using the `azure-key-vault-secret` k8s secret.\n",
    "\n",
    "The code tried to retrieve the 2 secrets described above from the key vault. Note that the `with_secrets()` function is called with an empty list of secrets, which will provide access to all the secrets that are accessible within the key vault. If you wish to limit access to specific keys, then modify the code to: `\"secrets\": [\"secret_name\",...]` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_name = \"azure-key-vault-project\"\n",
    "azure_key_vault_k8s_secret = \"azure-key-vault-secret\"\n",
    "azure_key_vault_name = \"azure-demo-key-vault\"\n",
    "\n",
    "task = mlrun.new_task(\n",
    "    project=proj_name,\n",
    "    name=\"azure_vault_test_run\",\n",
    "    handler=\"vault_func\",\n",
    "    out_path=\"/home/jovyan/examples\",\n",
    "    params={\"secrets\": [\"demo-secret-1\", \"demo-secret-2\"]},\n",
    ")\n",
    "\n",
    "task.with_secrets(\n",
    "    \"azure_vault\",\n",
    "    {\n",
    "        \"name\": azure_key_vault_name,\n",
    "        \"k8s_secret\": azure_key_vault_k8s_secret,\n",
    "        \"secrets\": [],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code executes the secret-printing function defined above with access to the Azure secrets specified in the task properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = func.run(task)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
